package chatglm

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/869413421/wechatbot/config"
	"github.com/dgrijalva/jwt-go"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"
)

const BASEURL = "https://open.bigmodel.cn/api/paas/v3/model-api/chatglm_std/sse-invoke"

// ChatGLMResponseBody 请求体
type ChatGLMResponseBody struct {
	ID    string                  `json:"id"`
	Event string                  `json:"event"`
	Data  string                  `json:"data"`
	Meta  ChatGLMResponseBodyMeta `json:"meta"`
}

type ChatGLMResponseBodyMeta struct {
	ChatGLMResponseBodyUsage ChatGLMResponseBodyUsage `json:"usage"`
}

type ChatGLMResponseBodyUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type ChoiceItem struct {
}

// ChatGLMRequestBody 响应体
type ChatGLMRequestBody struct {
	Prompt      ChatGLMPrompt `json:"prompt"`
	MaxTokens   int           `json:"max_tokens"`
	Temperature float32       `json:"temperature"`
	TopP        int           `json:"top_p"`
	RequestId   string        `json:"request_id"`
	Incremental bool          `json:"incremental"`
}

type ChatGLMPrompt struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// Completions chatglm文本模型回复
// curl https://open.bigmodel.cn/api/paas/v3/model-api/chatglm_std/sse-invoke
// -H "Content-Type: application/json"
// -H "Authorization: Bearer your chatGPT key"
// -d '{"model": "text-davinci-003", "prompt": "give me good song", "temperature": 0, "max_tokens": 7}'
func Completions(msg string, sender string) (string, error) {
	prompt := ChatGLMPrompt{
		Role:    "user",
		Content: msg,
	}
	requestBody := ChatGLMRequestBody{
		Prompt:      prompt,
		MaxTokens:   2048,
		Temperature: 0.7,
		TopP:        1,
		Incremental: true,
	}
	requestData, err := json.Marshal(requestBody)

	if err != nil {
		return "", err
	}
	log.Printf("request chatglm json string : %v", string(requestData))
	req, err := http.NewRequest("POST", BASEURL, bytes.NewBuffer(requestData))
	if err != nil {
		return "", err
	}

	apiKey := config.LoadConfig().ApiKey
	token, err := generateToken(apiKey)
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+token)
	client := &http.Client{}
	response, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()

	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}

	var reply string
	log.Println("body" + string(body))
	lines := strings.Split(string(body), "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "data:") {
			data := strings.TrimPrefix(line, "data:")
			reply += data
		}
	}

	log.Printf("chatglm response text: %s \n", reply)
	return reply, nil
}

func generateToken(apikey string) (string, error) {
	stringSlice := strings.Split(apikey, ".")
	id := ""
	secret := ""
	if len(stringSlice) > 1 {
		id = stringSlice[0]
		secret = stringSlice[1]
	}

	if id == "" || secret == "" {
		return "", errors.New("invalid apikey")
	}

	claims := jwt.MapClaims{}
	claims["api_key"] = id
	claims["timestamp"] = time.Now().UnixNano() / 1e6
	claims["exp"] = time.Now().Add(time.Minute*10).UnixNano() / 1e6

	mySigningKey := []byte(secret)
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	token.Header["alg"] = "HS256"
	token.Header["sign_type"] = "SIGN"

	signedToken, err := token.SignedString(mySigningKey)
	fmt.Println(signedToken, err)
	return signedToken, err
}
